Attention

多头缩放点积注意力机制（Scaled Dot-Product Attention）

\[\text{Attention}(Q, K, V) = \operatorname{softmax}\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V\]
输入：

  • Q - 查询矩阵地址（行优先），形状 \([B, H, L, D]\) 展平。

  • K - 键矩阵地址（行优先），形状 \([B, H, L, D]\) 展平。

  • V - 值矩阵地址（行优先），形状 \([B, H, L, D]\) 展平。

  • batch_size (B) - 批大小。

  • seq_len (L) - 序列长度。

  • head_num (H) - 多头数量。

  • head_dim (D) - 每头通道维数。

  • QK, softmax_out - 中间缓冲区地址，容量不小于 \(B\times H\times L\times L\)。

  • core_mask（可选） - 核掩码（仅适用于共享存储版本）。

输出：

  • output - 输出地址（行优先），形状 \([B, H, L, D]\) 展平。

支持平台：

FT78NE MT7004

备注

  • 当前实现基于 fp32；输入/中间/输出缓冲区不应重叠。

  • 内存布局为行优先（row-major）。

共享存储版本：

void fp_attention_s(float *Q, float *K, float *V, float *output, int batch_size, int seq_len, int head_num, int head_dim, float *QK, float *softmax_out, int core_mask)

C调用示例：

 1#include <stdio.h>
 2
 3int main(int argc, char* argv[]) {
 4    int B = 2, L = 128, H = 8, D = 64;
 5    float *Q = (float *)0xA0000000;      // DDR
 6    float *K = (float *)0xA1000000;      // DDR
 7    float *V = (float *)0xA2000000;      // DDR
 8    float *O = (float *)0xA3000000;      // DDR
 9    float *QK = (float *)0xA4000000;     // DDR
10    float *SM = (float *)0xA5000000;     // DDR
11    int core_mask = 0xff;
12    fp_attention_s(Q, K, V, O, B, L, H, D, QK, SM, core_mask);
13    return 0;
14}

私有存储版本：

void fp_attention_p(float *Q, float *K, float *V, float *output, int batch_size, int seq_len, int head_num, int head_dim, float *QK, float *softmax_out)

C调用示例：

 1#include <stdio.h>
 2
 3int main(int argc, char* argv[]) {
 4    int B = 1, L = 64, H = 4, D = 32;
 5    float *Q = (float *)0x10000000;   // L2
 6    float *K = (float *)0x10040000;   // L2
 7    float *V = (float *)0x10080000;   // L2
 8    float *O = (float *)0x100C0000;   // L2
 9    float *QK = (float *)0x10100000;  // L2
10    float *SM = (float *)0x10200000;  // L2
11    fp_attention_p(Q, K, V, O, B, L, H, D, QK, SM);
12    return 0;
13}